import numpy as np
import pandas as pd
import torch
import torch.nn as nn
from sklearn.preprocessing import StandardScaler, MinMaxScaler
from torch.utils.data import TensorDataset
import os

class Config():
    timestep = 12  # 时间步长，就是利用多少时间窗口
    batch_size = 32   # 批次大小
    feature_size = 1  # 每个步长对应的特征数量，这里只使用1维，每天的风速
    hidden_size = 256  # 隐层大小
    output_size = 1  # 由于是单输出任务，最终输出层大小为1，预测未来1天风速
    num_layers = 2  # gru的层数
    epochs = 200 # 迭代轮数
    best_loss = 0 # 记录损失
    learning_rate = 0.001 # 学习率
    model_name = 'gru' # 模型名称
    save_path = './{}.pth'.format(model_name) # 最优模型保存路径

config = Config()

class GRU(nn.Module):
    def __init__(self, feature_size, hidden_size, num_layers, output_size):
        super(GRU, self).__init__()
        self.hidden_size = hidden_size  # 隐层大小
        self.num_layers = num_layers  # gru层数
        # feature_size为特征维度，就是每个时间点对应的特征数量，这里为1
        self.gru = nn.GRU(feature_size, hidden_size, num_layers, batch_first=True)
        self.fc = nn.Linear(hidden_size, output_size)

    def forward(self, x, hidden=None):
        batch_size = x.shape[0]  # 获取批次大小

        # 初始化隐层状态
        if hidden is None:
            h_0 = x.data.new(self.num_layers, batch_size, self.hidden_size).fill_(0).float()
        else:
            h_0 = hidden

        # GRU运算
        output, h_0 = self.gru(x, h_0)

        # 获取GRU输出的维度信息
        batch_size, timestep, hidden_size = output.shape

        # 将output变成 batch_size * timestep, hidden_dim
        output = output.reshape(-1, hidden_size)

        # 全连接层
        output = self.fc(output)  # 形状为batch_size * timestep, 1

        # 转换维度，用于输出
        output = output.reshape(timestep, batch_size, -1)

        # 我们只需要返回最后一个时间片的数据即可
        return output[-1]

def split_data(data, timestep, feature_size):
    dataX = []  # 保存X
    dataY = []  # 保存Y

    # 将整个窗口的数据保存到X中，将未来一天保存到Y中
    for index in range(len(data) - timestep):
        dataX.append(data[index: index + timestep][:, 0])
        dataY.append(data[index + timestep][0])

    dataX = np.array(dataX)
    dataY = np.array(dataY)

    # 获取训练集大小
    train_size = 1000

    # 划分训练集、测试集
    x_train = dataX[: train_size, :].reshape(-1, timestep, feature_size)
    y_train = dataY[: train_size].reshape(-1, 1)

    x_test = dataX[train_size:, :].reshape(-1, timestep, feature_size)
    y_test = dataY[train_size:].reshape(-1, 1)

    return [x_train, y_train, x_test, y_test]

def test_to_tensor(x_test, y_test):
    x_test_tensor = torch.from_numpy(x_test).to(torch.float32)
    y_test_tensor = torch.from_numpy(y_test).to(torch.float32)
    test_data = TensorDataset(x_test_tensor, y_test_tensor)
    return x_test_tensor,y_test_tensor,test_data

'''
对外调用接口;
提供给外部其他Python文件调用接口.
'''

def data_process(data,url):
    if url:
        df = pd.read_excel(url)
    else:
        df = data
    df.columns = ['y']
    scaler = MinMaxScaler()
    df1 = scaler.fit_transform(np.array(df.values))
    x_train, y_train, x_test, y_test = split_data(df1, config.timestep, config.feature_size)
    x_test_tensor, y_test_tensor, test_data = test_to_tensor(x_test, y_test)

    # 当前程序路径问题，进行识别当前路径转换
    #curr_url = os.path.join(os.getcwd(), 'feature', 'GRU_model')
    # load the saved model
    model = GRU(feature_size=config.feature_size, hidden_size=config.hidden_size, num_layers=config.num_layers,
                output_size=config.output_size)
    # 当前程序路径问题，进行识别当前路径转换
    curr_url = os.path.join(os.getcwd(), 'nankai', 'GRU_model')
    # model.load_state_dict(torch.load('GRU_model'))
    model.load_state_dict(torch.load(curr_url))
    y_test_pred = scaler.inverse_transform(model(x_test_tensor).detach().numpy()[: 100]).tolist()
    return y_test_pred

if __name__ =="__main__":

    # 读取文件
    url = 'GRU_data.xlsx'
    data = pd.read_excel('GRU_data.xlsx')
    y_test_pred = data_process(data,url)

    # 输出预测值
    print('Predict:', y_test_pred)